﻿/*
 * @author: wizardc
 */

using System.Collections.Generic;
using Dou.Job;
using Unity.Burst;
using Unity.Collections;
using Unity.Collections.LowLevel.Unsafe;
using Unity.Jobs;
using Unity.Mathematics;
using UnityEngine;

namespace Dou.GPU
{
    /// <summary>
    /// 同一个批次的 GPU 动画对象
    /// </summary>
    public unsafe partial class GPUBatchAnimation
    {
        private static readonly int ID_TransformBuf = Shader.PropertyToID("transformBuf");
        private static readonly int ID_PlayIndexBuf = Shader.PropertyToID("playIndexBuf");
        private static int _uid = 0;

        private GPUInstancingMgr _gpuInst;
        
        private int _batchId;
        private GPUAnimData _animData;

        private Dictionary<int, GPUAnimation> _animationMap = new Dictionary<int, GPUAnimation>();
        
        public GPUBatchAnimation(int batchId, GPUAnimData animData)
        {
            _gpuInst = GPUInstancingMgr.instance;
            
            _batchId = batchId;
            _animData = animData;
            
            _gpuInst.RegisterProperty(_batchId, ID_TransformBuf, GPUDefine.Stride_Float4x4);
            _gpuInst.RegisterProperty(_batchId, ID_PlayIndexBuf, GPUDefine.Stride_Int);
        }

        internal IEnumerable<int> EnumerateUid()
        {
            foreach (var kv in _animationMap)
            {
                yield return kv.Key;
            }
        }

        public int CreateAnimation()
        {
            var animation = new GPUAnimation(_animData);
            _animationMap.Add(_uid, animation);
            return _uid++;
        }

        public GPUAnimation GetAnimation(int uid)
        {
            if (_animationMap.ContainsKey(uid))
            {
                return _animationMap[uid];
            }
            return null;
        }
        
        public void SetPosition(int uid, float3 position)
        {
            var animation = GetAnimation(uid);
            if (animation != null)
            {
                animation.data.position = position;
            }
        }

        public void SetRotation(int uid, quaternion rotation)
        {
            var animation = GetAnimation(uid);
            if (animation != null)
            {
                animation.data.rotation = rotation;
            }
        }

        public void SetScale(int uid, float3 scale)
        {
            var animation = GetAnimation(uid);
            if (animation != null)
            {
                animation.data.scale = scale;
            }
        }
        
        public void Play(int uid, string animName, float time = 0)
        {
            var animation = GetAnimation(uid);
            animation?.Play(animName, time);
        }
        
        public void SetIsLoop(int uid, bool isLoop)
        {
            var animation = GetAnimation(uid);
            animation?.SetIsLoop(isLoop);
        }

        public void SetTimeScale(int uid, float timeScale)
        {
            var animation = GetAnimation(uid);
            animation?.SetTimeScale(timeScale);
        }

        public void Pause(int uid)
        {
            var animation = GetAnimation(uid);
            animation?.Pause();
        }

        public void Resume(int uid)
        {
            var animation = GetAnimation(uid);
            animation?.Resume();
        }

        public void Stop(int uid)
        {
            var animation = GetAnimation(uid);
            animation?.Stop();
        }
        
        [BurstCompile]
        public void Update(float deltaTime)
        {
            var count = _animationMap.Count;
            var dataArray = new NativeArray<GPUAnimationData>(count, Allocator.TempJob);
            var index2UidArray = new NativeArray<int>(count, Allocator.TempJob);
            var index = 0;
            foreach (var kv in _animationMap)
            {
                var animation = kv.Value;
                dataArray[index] = animation.data;
                index2UidArray[index] = kv.Key;
                ++index;
            }
            var transArray = new NativeArray<float4x4>(count, Allocator.TempJob);
            var indexArray = new NativeArray<int>(count, Allocator.TempJob);

            var dataPtr = (GPUAnimationData*) dataArray.GetUnsafePtr();
            var index2UidPtr = (int*) index2UidArray.GetUnsafePtr();
            var transPtr = (float4x4*) transArray.GetUnsafePtr();
            var indexPtr = (int*) indexArray.GetUnsafePtr();

            var managedUpdateJob = new GPUAnimationUpdateJob();
            var updateJob = new ManagedJobParallelFor<GPUAnimationUpdateJob>();
            updateJob.Work = managedUpdateJob;

            managedUpdateJob.dataPtr = dataPtr;
            managedUpdateJob.index2UidPtr = index2UidPtr;
            managedUpdateJob.transPtr = transPtr;
            managedUpdateJob.indexPtr = indexPtr;
            managedUpdateJob.deltaTime = deltaTime;
            managedUpdateJob.animationMap = _animationMap;
            
            var handle = updateJob.Schedule(count, 64);
            handle.Complete();
            JobHandle.ScheduleBatchedJobs();
            
            GPUInstancingMgr.instance.SetFloat4x4(_batchId, ID_TransformBuf, ref transArray);
            GPUInstancingMgr.instance.SetInt(_batchId, ID_PlayIndexBuf, ref indexArray);
            
            dataArray.Dispose();
            index2UidArray.Dispose();
            transArray.Dispose();
            indexArray.Dispose();
        }

        public bool DestroyAnimation(int uid)
        {
            if (!_animationMap.ContainsKey(uid))
            {
                return false;
            }
            _animationMap.Remove(uid);
            return true;
        }

        public void Dispose()
        {
            GPUInstancingMgr.instance.UnregisterBatch(_batchId);
            _animationMap.Clear();
        }
    }

    [BurstCompile]
    [Il2CppSetOption(Option.ArrayBoundsChecks, false)]
    [Il2CppSetOption(Option.NullChecks, false)]
    public unsafe class GPUAnimationUpdateJob : ManagedJobParallelFor<GPUAnimationUpdateJob>.IWork
    {
        public GPUAnimationData* dataPtr;
        public int* index2UidPtr;
        public float4x4* transPtr;
        public int* indexPtr;
        
        public Dictionary<int, GPUAnimation> animationMap;
        public float deltaTime;
        
        [Il2CppSetOption(Option.ArrayBoundsChecks, false)]
        [Il2CppSetOption(Option.NullChecks, false)]
        public void Execute(int index)
        {
            var ptr = dataPtr + index;
            
            if (ptr->isPlaying)
            {
                ptr->playTime += deltaTime * ptr->timeScale;
                if (ptr->playTime > ptr->totalTime)
                {
                    if (ptr->isLoop)
                    {
                        ptr->playTime %= ptr->totalTime;
                    }
                    else
                    {
                        ptr->playTime = ptr->totalTime;
                    }
                }
            }
            ptr->playIndex = ptr->startIndex + (int) math.floor(ptr->playTime / ptr->frameTime);

            var animation = animationMap[*(index2UidPtr + index)];
            animation.data.playTime = ptr->playTime;
            animation.data.playIndex = ptr->playIndex;
            
            *(transPtr + index) = float4x4.TRS(ptr->position, ptr->rotation, ptr->scale);
            *(indexPtr + index) = ptr->playIndex;
        }
    }
}
